{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l2mMvIUG9meX"
   },
   "source": [
    "# Debugging Model Training Performance using Profiler\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dmfmQL6w84pS"
   },
   "source": [
    "## Learning objectives\n",
    "\n",
    "1. Create a Cloud Storage bucket.\n",
    "2. Create a TensorBoard instance.\n",
    "3. Create and run a custom training job.\n",
    "4. View the TensorBoard Profiler dashboard\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tvgnzT1CKxrO"
   },
   "source": [
    "## Overview\n",
    "\n",
    "In this notebook, you will learn how to enable Vertex AI TensorBoard Profiler so you can debug model training performance for your custom training jobs.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zfXf0r-K81Y-"
   },
   "source": [
    "### Dataset\n",
    "\n",
    "The dataset used for this tutorial is the [mnist dataset](https://www.tensorflow.org/datasets/catalog/mnist) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview).\n",
    "\n",
    "Each learning objective will correspond to a __#TODO__ in the notebook, where you will complete the notebook cell's code before running the cell. Refer to the [solution notebook](../solutions/custom_training_tensorboard_profiler.ipynb) for reference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i7EUnXsZhAGF"
   },
   "source": [
    "### Install required packages\n",
    "\n",
    "Install the following packages required to execute this notebook. \n",
    "\n",
    "- Ensure that you're using Tensorflow 2.4 or a later version.\n",
    "- Install the Vertex AI SDK with the cloud-profiler plugin. From your local Docker container, run `pip install google-cloud-aiplatform[cloud_profiler]`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "2b4ef9b72d43"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# The Vertex AI Workbench Notebook product has specific requirements\n",
    "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
    "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
    "    \"/opt/deeplearning/metadata/env_version\"\n",
    ")\n",
    "\n",
    "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
    "USER_FLAG = \"\"\n",
    "if IS_WORKBENCH_NOTEBOOK:\n",
    "    USER_FLAG = \"--user\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "Asm0sNIRjz6b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33m  WARNING: The scripts f2py, f2py3 and f2py3.7 are installed in '/home/jupyter/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33m  WARNING: The script tensorboard is installed in '/home/jupyter/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33m  WARNING: The scripts estimator_ckpt_converter, import_pb_to_tensorboard, saved_model_cli, tensorboard, tf_upgrade_v2, tflite_convert, toco and toco_from_protos are installed in '/home/jupyter/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "tfx-bsl 1.10.1 requires google-api-python-client<2,>=1.7.11, but you have google-api-python-client 2.63.0 which is incompatible.\n",
      "tfx-bsl 1.10.1 requires pyarrow<7,>=6, but you have pyarrow 9.0.0 which is incompatible.\n",
      "tensorflow-transform 1.10.1 requires pyarrow<7,>=6, but you have pyarrow 9.0.0 which is incompatible.\n",
      "tensorflow-transform 1.10.1 requires tensorflow!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,!=2.8.*,<2.10,>=1.15.5, but you have tensorflow 2.10.0 which is incompatible.\n",
      "tensorflow-io 0.21.0 requires tensorflow<2.7.0,>=2.6.0, but you have tensorflow 2.10.0 which is incompatible.\n",
      "tensorflow-io 0.21.0 requires tensorflow-io-gcs-filesystem==0.21.0, but you have tensorflow-io-gcs-filesystem 0.27.0 which is incompatible.\n",
      "pandas-profiling 3.0.0 requires tangled-up-in-unicode==0.1.0, but you have tangled-up-in-unicode 0.2.0 which is incompatible.\n",
      "apache-beam 2.41.0 requires dill<0.3.2,>=0.3.1.1, but you have dill 0.3.5.1 which is incompatible.\n",
      "apache-beam 2.41.0 requires pyarrow<8.0.0,>=0.15.1, but you have pyarrow 9.0.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m\u001b[33m  WARNING: The script tb-gcp-uploader is installed in '/home/jupyter/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "! pip3 install --upgrade tensorflow {USER_FLAG} -q\n",
    "! pip3 install --upgrade google-cloud-aiplatform[cloud_profiler] {USER_FLAG} -q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hhq5zEbGg0XX"
   },
   "source": [
    "### Restart the kernel\n",
    "\n",
    "After you install the required packages, you need to restart the notebook kernel so it can find the packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "EzrelQZ22IZj"
   },
   "outputs": [],
   "source": [
    "# Automatically restart kernel after installs\n",
    "import os\n",
    "\n",
    "if not os.getenv(\"IS_TESTING\"):\n",
    "    # Automatically restart kernel after installs\n",
    "    import IPython\n",
    "\n",
    "    app = IPython.Application.instance()\n",
    "    app.kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lWEdiXsJg0XY"
   },
   "source": [
    "## Before you begin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WReHDGG5g0XY"
   },
   "source": [
    "### Set your project ID\n",
    "\n",
    "**If you don't know your project ID**, you can get a list of projects available to you using `gcloud`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "oM1iC_MfAts1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Project ID:  qwiklabs-gcp-00-1fe7a1835a07\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "PROJECT_ID = \"[your-project-id]\" # Replace with your Project ID\n",
    "\n",
    "# Get your Google Cloud project ID from gcloud\n",
    "if not os.getenv(\"IS_TESTING\"):\n",
    "    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null\n",
    "    PROJECT_ID = shell_output[0]\n",
    "    print(\"Project ID: \", PROJECT_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "set_gcloud_project_id"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updated property [core/project].\n"
     ]
    }
   ],
   "source": [
    "! gcloud config set project $PROJECT_ID"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JWHp43-vHVgL"
   },
   "source": [
    "### Get your project number\n",
    "\n",
    "Now that the project ID is set, you get your corresponding project number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "yMAT0gATHe1V"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Project number: 21417957250\n"
     ]
    }
   ],
   "source": [
    "shell_output = ! gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
    "PROJECT_NUMBER = shell_output[0]\n",
    "print(\"Project number:\", PROJECT_NUMBER)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "region"
   },
   "source": [
    "### Region\n",
    "\n",
    "You can also change the `REGION` variable, which is used for operations\n",
    "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
    "\n",
    "- Americas: `us-central1`\n",
    "- Europe: `europe-west4`\n",
    "- Asia Pacific: `asia-east1`\n",
    "\n",
    "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
    "\n",
    "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "kmPv1dw33srh"
   },
   "outputs": [],
   "source": [
    "REGION = \"[your-region]\"  # Replace with your region\n",
    "\n",
    "if REGION == \"[your-region]\":\n",
    "    REGION = \"us-central1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "06571eb4063b"
   },
   "source": [
    "#### Timestamp\n",
    "\n",
    "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append it onto the name of resources you create in this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "697568e92bd6"
   },
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M8jPMDYIL4kc"
   },
   "source": [
    "#### Enable the Vertex AI API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "2T-khaOxMBql"
   },
   "outputs": [],
   "source": [
    "! gcloud services enable aiplatform.googleapis.com"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zgPO1eR3CYjk"
   },
   "source": [
    "### Create a Cloud Storage bucket\n",
    "\n",
    "**The following steps are required, regardless of your notebook environment.**\n",
    "\n",
    "When you submit a training job using the Cloud SDK, you upload a Python package containing your training code to a Cloud Storage bucket. Vertex AI runs the code from this package. The Cloud Storage bucket is also used to store TensorBoard logs generated by your training jobs.\n",
    "\n",
    "Set the name of your Cloud Storage bucket below. It must be unique across all Cloud Storage buckets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "MzGDU7TWdts_"
   },
   "outputs": [],
   "source": [
    "BUCKET_NAME = \"qwiklabs-gcp-00-1fe7a1835a07\"  # Replace with your unique bucket name\n",
    "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-EcIXiGsCePi"
   },
   "source": [
    "If you don't have a Cloud Storage bucket, create one by running the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "NIq7R4HZCfIc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating gs://qwiklabs-gcp-00-1fe7a1835a07/...\n"
     ]
    }
   ],
   "source": [
    "# TODO 1\n",
    "# Create a regional Cloud Storage bucket\n",
    "# TODO 1: Your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ucvCsknMCims"
   },
   "source": [
    "Finally, validate access to your Cloud Storage bucket by examining its contents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "vhOb7YnwClBb"
   },
   "outputs": [],
   "source": [
    "! gcloud storage ls --all-versions --long $BUCKET_URI"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "set_service_account"
   },
   "source": [
    "### Set up your service account\n",
    "\n",
    "The Vertex AI TensorBoard integration with custom training requires attaching a service account with the Storage Admin role (roles/storage.admin) and Vertex AI User role (roles/aiplatform.user)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cZtfI9tKATVT"
   },
   "source": [
    "#### Set your service account\n",
    "\n",
    "If you do not want to use your project's Compute Engine service account, set `SERVICE_ACCOUNT` to another service account ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "id": "vYE3b942wza4"
   },
   "outputs": [],
   "source": [
    "SERVICE_ACCOUNT = \"21417957250-compute@developer.gserviceaccount.com\"  # Replace '[your-service-account]' with your project's Compute Engine service account"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jmYVMGNOJHVo"
   },
   "source": [
    "#### Set service account access for Vertex AI TensorBoard integration\n",
    "\n",
    "Run the following commands to grant your service account access to the Storage Admin role (roles/storage.admin) and Vertex AI User role (roles/aiplatform.user)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "LP8jEAt8xrtn"
   },
   "outputs": [],
   "source": [
    "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/aiplatform.user\n",
    "\n",
    "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XoEqT2Y4DJmf"
   },
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "pRUOFELefqf1"
   },
   "outputs": [],
   "source": [
    "import google.cloud.aiplatform as aiplatform"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "init_aip:mbsdk,all"
   },
   "source": [
    "### Initialize Vertex AI SDK for Python\n",
    "\n",
    "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "cChm5w3j3srj"
   },
   "outputs": [],
   "source": [
    "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xHUW20OnjrEw"
   },
   "source": [
    "#### Set hardware accelerators\n",
    "\n",
    "You can set hardware accelerators for training.\n",
    "\n",
    "Set the variables `TRAIN_GPU/TRAIN_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Tesla K80 GPUs allocated to each VM, you would specify:\n",
    "\n",
    "    (aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
    "\n",
    "See the [locations where accelerators are available](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators).\n",
    "\n",
    "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "52UoZKMBjx3I"
   },
   "outputs": [],
   "source": [
    "if os.getenv(\"IS_TESTING_TRAIN_GPU\"):\n",
    "    TRAIN_GPU, TRAIN_NGPU = (\n",
    "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
    "        int(os.getenv(\"IS_TESTING_TRAIN_GPU\")),\n",
    "    )\n",
    "else:\n",
    "    TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
    "\n",
    "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
    "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
    "        aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
    "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
    "    )\n",
    "else:\n",
    "    DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q4wlRWnJj4aR"
   },
   "source": [
    "#### Set pre-built containers\n",
    "\n",
    "\n",
    "Set the pre-built Docker container image for training and prediction.\n",
    "\n",
    "\n",
    "For the latest list, see [Pre-built containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers).\n",
    "\n",
    "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "8SYq7K0VcYTg"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training image: us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-9:latest None None\n",
      "Deployment image: us-docker.pkg.dev/vertex-ai/training/tf2-cpu.2-9:latest None None\n"
     ]
    }
   ],
   "source": [
    "TF = \"2-9\"\n",
    "\n",
    "if TRAIN_GPU:\n",
    "    TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
    "else:\n",
    "    TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
    "\n",
    "if DEPLOY_GPU:\n",
    "    DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
    "else:\n",
    "    DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
    "\n",
    "TRAIN_IMAGE = \"us-docker.pkg.dev/vertex-ai/training/{}:latest\".format(TRAIN_VERSION)\n",
    "DEPLOY_IMAGE = \"us-docker.pkg.dev/vertex-ai/training/{}:latest\".format(DEPLOY_VERSION)\n",
    "\n",
    "print(\"Training image:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
    "print(\"Deployment image:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-ayTbNdi62_t"
   },
   "source": [
    "## Create a TensorBoard instance\n",
    "\n",
    "A Vertex AI TensorBoard instance, which is a regionalized resource storing your Vertex AI TensorBoard experiments, must be created before the experiments can be visualized. You can create multiple instances in a project. You can use command  `gcloud ai tensorboards list` to get a list of your existing TensorBoard instances."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9c3QrDTZdaxk"
   },
   "source": [
    "#### Set your TensorBoard instance display name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "azlwb__AX8gs"
   },
   "outputs": [],
   "source": [
    "TENSORBOARD_NAME = \"tensorboard-tb\"  # Replace with any name of your choice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vJrWKK0mY7H7"
   },
   "source": [
    "If you don't have a TensorBoard instance, create one by running the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "JqVNsRFrc_78"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating Tensorboard\n",
      "Create Tensorboard backing LRO: projects/21417957250/locations/us-central1/tensorboards/5219117814262005760/operations/6713704104704606208\n",
      "Tensorboard created. Resource name: projects/21417957250/locations/us-central1/tensorboards/5219117814262005760\n",
      "To use this Tensorboard in another session:\n",
      "tb = aiplatform.Tensorboard('projects/21417957250/locations/us-central1/tensorboards/5219117814262005760')\n"
     ]
    }
   ],
   "source": [
    "# TODO 2\n",
    "tensorboard = # TODO 2: Your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6dECoKEBcsfu"
   },
   "source": [
    "Once your TensorBoard instance is created, you get the TensorBoard instance name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "m8Rf-qDIegdU"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorBoard display name: tensorboard-tb\n",
      "TensorBoard instance name: projects/21417957250/locations/us-central1/tensorboards/5219117814262005760\n"
     ]
    }
   ],
   "source": [
    "TENSORBOARD_INSTANCE_NAME = tensorboard.resource_name\n",
    "\n",
    "print(\"TensorBoard display name:\", TENSORBOARD_NAME)\n",
    "print(\"TensorBoard instance name:\", TENSORBOARD_INSTANCE_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yoR29gW2S24w"
   },
   "source": [
    "## Train a model\n",
    "\n",
    "To train a model using your custom training code, choose one of the following options:\n",
    "\n",
    "- **Prebuilt container**: Load your custom training code as a Python package to a prebuilt container image from Google Cloud.\n",
    "\n",
    "- **Custom container**: Create your own container image that contains your custom training code.\n",
    "\n",
    "In this tutorial, we will train a custom model using a Google Cloud prebuild container."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TEbNcvYwmDRH"
   },
   "source": [
    "### Examine the training package\n",
    "\n",
    "#### Package layout\n",
    "\n",
    "Before you start the training, let's take a look at how a Python package is assembled for a custom training job. When extracted, the package contains the following:\n",
    "\n",
    "- PKG-INFO\n",
    "- README.md\n",
    "- setup.cfg\n",
    "- setup.py\n",
    "- trainer\n",
    "  - \\_\\_init\\_\\_.py\n",
    "  - task.py\n",
    "\n",
    "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the docker image.\n",
    "\n",
    "The file `trainer/task.py` is the Python script for executing the custom training job. Notice that when referring to this file in worker_pool_spec, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "D_dtjarIRuMs"
   },
   "outputs": [],
   "source": [
    "PYTHON_PACKAGE_APPLICATION_DIR = \"custom\"\n",
    "\n",
    "source_package_file_name = f\"{PYTHON_PACKAGE_APPLICATION_DIR}/dist/trainer-0.1.tar.gz\"\n",
    "python_package_gcs_uri = f\"{BUCKET_URI}/trainer-0.1.tar.gz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "aDilgzPGeyKs"
   },
   "outputs": [],
   "source": [
    "# Make folder for Python training script\n",
    "! rm -rf {PYTHON_PACKAGE_APPLICATION_DIR}\n",
    "! mkdir {PYTHON_PACKAGE_APPLICATION_DIR}\n",
    "\n",
    "# Add package information\n",
    "! touch {PYTHON_PACKAGE_APPLICATION_DIR}/README.md\n",
    "\n",
    "# Make the training subfolder\n",
    "! mkdir {PYTHON_PACKAGE_APPLICATION_DIR}/trainer\n",
    "! touch {PYTHON_PACKAGE_APPLICATION_DIR}/trainer/__init__.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "qiL565golGOF"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing ./custom/setup.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./{PYTHON_PACKAGE_APPLICATION_DIR}/setup.py\n",
    "\n",
    "from setuptools import find_packages\n",
    "from setuptools import setup\n",
    "import setuptools\n",
    "\n",
    "from distutils.command.build import build as _build\n",
    "import subprocess\n",
    "\n",
    "REQUIRED_PACKAGES = [\n",
    "    'google-cloud-aiplatform[cloud_profiler]',\n",
    "]\n",
    "\n",
    "setup(\n",
    "    install_requires=REQUIRED_PACKAGES,\n",
    "    packages=find_packages(),\n",
    "    include_package_data=True,\n",
    "    name='trainer',\n",
    "    version='0.1',\n",
    "    url=\"wwww.google.com\",\n",
    "    description='Vertex AI | Training | Python Package'\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hyAwgsoQmaYI"
   },
   "source": [
    "### Prepare the training script\n",
    "\n",
    "Your training code must be configured to write TensorBoard logs to a Cloud Storage bucket, the location of which Vertex AI Training automatically makes available through a predefined environment variable, `AIP_TENSORBOARD_LOG_DIR`.\n",
    "\n",
    "This can usually be done by providing `os.environ['AIP_TENSORBOARD_LOG_DIR']` as the log directory to the open source TensorBoard log writing APIs. \n",
    "\n",
    "For example, in TensorFlow 2.x, you can use following code to create a tensorboard_callback: \n",
    "\n",
    "    tensorboard_callback = tf.keras.callbacks.TensorBoard( \n",
    "      log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'], \n",
    "      histogram_freq=1) \n",
    "`AIP_TENSORBOARD_LOG_DIR` is in the `BASE_OUTPUT_DIR` that you provide when creating the custom training job.\n",
    "\n",
    "To enable Vertex AI TensorBoard Profiler for your training job, add the following to your training script:\n",
    "\n",
    "Add the cloud_profiler import at your top level imports:\n",
    "\n",
    "    from google.cloud.aiplatform.training_utils import cloud_profiler\n",
    "\n",
    "\n",
    "Initialize the cloud_profiler plugin by adding:\n",
    "\n",
    "\n",
    "    cloud_profiler.init()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "8JCgWW7Au1w8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing ./custom/trainer/task.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./{PYTHON_PACKAGE_APPLICATION_DIR}/trainer/task.py\n",
    "#!/usr/bin/env python\n",
    "\n",
    "import tensorflow as tf\n",
    "import argparse\n",
    "import os\n",
    "from google.cloud.aiplatform.training_utils import cloud_profiler\n",
    "import time\n",
    "\n",
    "\"\"\"Train an mnist model and use cloud_profiler for profiling.\"\"\"\n",
    "\n",
    "def _create_model():\n",
    "    model = tf.keras.models.Sequential(\n",
    "        [\n",
    "            tf.keras.layers.Flatten(input_shape=(28, 28)),\n",
    "            tf.keras.layers.Dense(128, activation=\"relu\"),\n",
    "            tf.keras.layers.Dropout(0.2),\n",
    "            tf.keras.layers.Dense(10),\n",
    "        ]\n",
    "    )\n",
    "    return model\n",
    "\n",
    "\n",
    "def main(args):\n",
    "    strategy = None\n",
    "    if args.distributed:\n",
    "        strategy = tf.distribute.MultiWorkerMirroredStrategy()\n",
    "\n",
    "    mnist = tf.keras.datasets.mnist\n",
    "\n",
    "    (x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "    x_train, x_test = x_train / 255.0, x_test / 255.0\n",
    "\n",
    "    if args.distributed:\n",
    "        strategy = tf.distribute.MultiWorkerMirroredStrategy()\n",
    "        with strategy.scope():\n",
    "            model = _create_model()\n",
    "            model.compile(\n",
    "                optimizer=\"adam\",\n",
    "                loss=tf.keras.losses.sparse_categorical_crossentropy,\n",
    "                metrics=[\"accuracy\"],\n",
    "            )\n",
    "    else:\n",
    "        model = _create_model()\n",
    "        model.compile(\n",
    "            optimizer=\"adam\",\n",
    "            loss=tf.keras.losses.sparse_categorical_crossentropy,\n",
    "            metrics=[\"accuracy\"],\n",
    "        )\n",
    "\n",
    "    # Initialize the profiler.\n",
    "    cloud_profiler.init()\n",
    "\n",
    "    # Use AIP_TENSORBOARD_LOG_DIR to update where logs are written to.\n",
    "    tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
    "       log_dir=os.environ[\"AIP_TENSORBOARD_LOG_DIR\"], histogram_freq=1\n",
    "    )\n",
    "\n",
    "    model.fit(\n",
    "        x_train,\n",
    "        y_train,\n",
    "        epochs=args.epochs,\n",
    "        verbose=0,\n",
    "        callbacks=[tensorboard_callback],\n",
    "    )\n",
    "    \n",
    "    MODEL_DIR = os.getenv(\"AIP_MODEL_DIR\")\n",
    "    \n",
    "    model.save(MODEL_DIR)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\n",
    "        \"--epochs\", type=int, default=100, help=\"Number of epochs to run model.\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--distributed\", action=\"store_true\", help=\"Use MultiWorkerMirroredStrategy\"\n",
    "    )\n",
    "    args = parser.parse_args()\n",
    "    main(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6z3B_i3sTuZp"
   },
   "source": [
    "Run the following command to create a source distribution, dist/trainer-0.1.tar.gz:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "4kMkJmf-OFKA"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running sdist\n",
      "running egg_info\n",
      "creating trainer.egg-info\n",
      "writing trainer.egg-info/PKG-INFO\n",
      "writing dependency_links to trainer.egg-info/dependency_links.txt\n",
      "writing requirements to trainer.egg-info/requires.txt\n",
      "writing top-level names to trainer.egg-info/top_level.txt\n",
      "writing manifest file 'trainer.egg-info/SOURCES.txt'\n",
      "reading manifest file 'trainer.egg-info/SOURCES.txt'\n",
      "writing manifest file 'trainer.egg-info/SOURCES.txt'\n",
      "running check\n",
      "warning: check: missing meta-data: either (author and author_email) or (maintainer and maintainer_email) must be supplied\n",
      "\n",
      "creating trainer-0.1\n",
      "creating trainer-0.1/trainer\n",
      "creating trainer-0.1/trainer.egg-info\n",
      "copying files to trainer-0.1...\n",
      "copying README.md -> trainer-0.1\n",
      "copying setup.py -> trainer-0.1\n",
      "copying trainer/__init__.py -> trainer-0.1/trainer\n",
      "copying trainer/task.py -> trainer-0.1/trainer\n",
      "copying trainer.egg-info/PKG-INFO -> trainer-0.1/trainer.egg-info\n",
      "copying trainer.egg-info/SOURCES.txt -> trainer-0.1/trainer.egg-info\n",
      "copying trainer.egg-info/dependency_links.txt -> trainer-0.1/trainer.egg-info\n",
      "copying trainer.egg-info/requires.txt -> trainer-0.1/trainer.egg-info\n",
      "copying trainer.egg-info/top_level.txt -> trainer-0.1/trainer.egg-info\n",
      "Writing trainer-0.1/setup.cfg\n",
      "creating dist\n",
      "Creating tar archive\n",
      "removing 'trainer-0.1' (and everything under it)\n"
     ]
    }
   ],
   "source": [
    "!cd {PYTHON_PACKAGE_APPLICATION_DIR} && python3 setup.py sdist --formats=gztar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W_tXf3fSmijU"
   },
   "source": [
    "Now upload the source distribution with training application to Cloud Storage bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "0LUV52d9OL-j"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copying file://custom/dist/trainer-0.1.tar.gz [Content-Type=application/x-tar]...\n",
      "/ [1 files][  1.7 KiB/  1.7 KiB]                                                \n",
      "Operation completed over 1 objects/1.7 KiB.                                      \n"
     ]
    }
   ],
   "source": [
    "# TODO 3\n",
    "# Copy the source distribution to Cloud Storage bucket\n",
    "# TODO 3: Your code goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QJBlrqAr0Mbe"
   },
   "source": [
    "Validate the source distribution exists on Cloud Storage bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "CIBCwumT0LJm"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      1729  2022-10-14T10:22:14Z  gs://qwiklabs-gcp-00-1fe7a1835a07/trainer-0.1.tar.gz\n",
      "TOTAL: 1 objects, 1729 bytes (1.69 KiB)\n"
     ]
    }
   ],
   "source": [
    "!gcloud storage ls --long {python_package_gcs_uri}"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k4e6OYmimqTR"
   },
   "source": [
    "## Create and run the custom training job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RfxbuzdVMjar"
   },
   "source": [
    "Configure a [custom job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) with the [pre-built container](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers) image for training code packaged as Python source distribution. \n",
    "\n",
    "**NOTE:** When using Vertex AI SDK for Python for submitting a training job, it creates a [training pipeline](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline) which launches the custom job on Vertex AI Training service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "jd6P-lQhFBPg"
   },
   "outputs": [],
   "source": [
    "JOB_NAME = \"custom_job_\" + TIMESTAMP\n",
    "MACHINE_TYPE = \"n1-standard-4\"\n",
    "\n",
    "base_output_dir = \"{}/{}\".format(BUCKET_URI, JOB_NAME)\n",
    "python_module_name = \"trainer.task\"\n",
    "\n",
    "EPOCHS = 20\n",
    "training_args = [\n",
    "    \"--epochs=\" + str(EPOCHS),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "GeLtnUSgxs4Z"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JOB_NAME=custom_job_20221014101628\n",
      "python_module_name=trainer.task\n",
      "PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI=us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-9:latest\n",
      "python_package_gcs_uri=gs://qwiklabs-gcp-00-1fe7a1835a07/trainer-0.1.tar.gz\n",
      "base_output_dir=gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628\n",
      "tensorboard=projects/21417957250/locations/us-central1/tensorboards/5219117814262005760\n",
      "service_account=21417957250-compute@developer.gserviceaccount.com\n"
     ]
    }
   ],
   "source": [
    "print(f\"JOB_NAME={JOB_NAME}\")\n",
    "print(f\"python_module_name={python_module_name}\")\n",
    "print(f\"PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI={TRAIN_IMAGE}\")\n",
    "print(f\"python_package_gcs_uri={python_package_gcs_uri}\")\n",
    "print(f\"base_output_dir={base_output_dir}\")\n",
    "print(f\"tensorboard={TENSORBOARD_INSTANCE_NAME}\")\n",
    "print(f\"service_account={SERVICE_ACCOUNT}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "id": "15EKAqYbCBmn"
   },
   "outputs": [],
   "source": [
    "# TODO 4\n",
    "# Create a custom training job\n",
    "job = aiplatform.CustomPythonPackageTrainingJob(\n",
    "# TODO 4: Your code goes here\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "vk4hHJz1_7OT"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Output directory:\n",
      "gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628 \n",
      "View Training:\n",
      "https://console.cloud.google.com/ai/platform/locations/us-central1/training/7902914372110909440?project=21417957250\n",
      "CustomPythonPackageTrainingJob projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440 current state:\n",
      "PipelineState.PIPELINE_STATE_PENDING\n",
      "CustomPythonPackageTrainingJob projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "View backing custom job:\n",
      "https://console.cloud.google.com/ai/platform/locations/us-central1/training/7968119809684537344?project=21417957250\n",
      "View tensorboard:\n",
      "https://us-central1.tensorboard.googleusercontent.com/experiment/projects+21417957250+locations+us-central1+tensorboards+5219117814262005760+experiments+7968119809684537344\n",
      "CustomPythonPackageTrainingJob projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomPythonPackageTrainingJob projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomPythonPackageTrainingJob projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomPythonPackageTrainingJob projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomPythonPackageTrainingJob run completed. Resource name: projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440\n",
      "Training did not produce a Managed Model returning None. Training Pipeline projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440 is not configured to upload a Model. Create the Training Pipeline with model_serving_container_image_uri and model_display_name passed in. Ensure that your training script saves to model to os.environ['AIP_MODEL_DIR'].\n"
     ]
    }
   ],
   "source": [
    "if TRAIN_GPU:\n",
    "    job.run(\n",
    "        args=training_args,\n",
    "        replica_count=1,\n",
    "        machine_type=MACHINE_TYPE,\n",
    "        accelerator_type=TRAIN_GPU.name,\n",
    "        accelerator_count=TRAIN_NGPU,\n",
    "        base_output_dir=base_output_dir,\n",
    "        tensorboard=TENSORBOARD_INSTANCE_NAME,\n",
    "        service_account=SERVICE_ACCOUNT,\n",
    "        sync=True,\n",
    "    )\n",
    "else:\n",
    "    job.run(\n",
    "        args=training_args,\n",
    "        replica_count=1,\n",
    "        machine_type=MACHINE_TYPE,\n",
    "        base_output_dir=base_output_dir,\n",
    "        tensorboard=TENSORBOARD_INSTANCE_NAME,\n",
    "        service_account=SERVICE_ACCOUNT,\n",
    "        sync=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JkEe2Nb_85UD"
   },
   "source": [
    "## View the TensorBoard Profiler dashboard\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0cE5UCYgdRxx"
   },
   "source": [
    "When the custom job state switches to `Running`, you can access the Vertex AI TensorBoard Profiler dashboard through the Custom jobs page or the Experiments page on the Google Cloud console. \n",
    "\n",
    "The Google Cloud guide to [Profile model training performance using Profiler](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-profiler) provides detailed instructions for accessing the Vertex AI TensorBoard Profiler dashboard and capturing a profiling session. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TpV-iwP9qw9c"
   },
   "source": [
    "## Cleaning up\n",
    "\n",
    "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
    "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
    "\n",
    "Otherwise, you can delete the individual resources you created in this tutorial:\n",
    "\n",
    "- Training job\n",
    "- TensorBoard instance\n",
    "- Cloud Storage bucket\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "id": "WR-ZhQ9XwpRI"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deleting CustomPythonPackageTrainingJob : projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440\n",
      "Delete CustomPythonPackageTrainingJob  backing LRO: projects/21417957250/locations/us-central1/operations/7150553268559544320\n",
      "CustomPythonPackageTrainingJob deleted. . Resource name: projects/21417957250/locations/us-central1/trainingPipelines/7902914372110909440\n",
      "Deleting Tensorboard : projects/21417957250/locations/us-central1/tensorboards/5219117814262005760\n",
      "Delete Tensorboard  backing LRO: projects/21417957250/locations/us-central1/operations/4776030365028450304\n",
      "Tensorboard deleted. . Resource name: projects/21417957250/locations/us-central1/tensorboards/5219117814262005760\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/trainer-0.1.tar.gz#1665742934265426...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/logs/train/#1665743342029476...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/#1665743341627028...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/model/variables/variables.data-00000-of-00001#1665743423670416...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/model/variables/variables.index#1665743423969079...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/model/assets/#1665743424749710...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/model/keras_metadata.pb#1665743425321147...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/logs/#1665743341831803...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/model/saved_model.pb#1665743425107207...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/logs/train/events.out.tfevents.1665743342.3cc52d2a7217.371.0.v2#1665743420196163...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/model/variables/#1665743421877873...\n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/custom_job_20221014101628/model/#1665743421640137...\n",
      "/ [12/12 objects] 100% Done                                                     \n",
      "Operation completed over 12 objects.                                             \n",
      "Removing gs://qwiklabs-gcp-00-1fe7a1835a07/...\n"
     ]
    }
   ],
   "source": [
    "delete_customjob = True\n",
    "delete_tensorboard = True\n",
    "delete_bucket = True\n",
    "\n",
    "try:\n",
    "    if delete_customjob:\n",
    "        job.delete()\n",
    "except Exception as e:\n",
    "    print(e)\n",
    "\n",
    "try:\n",
    "    if delete_tensorboard:\n",
    "        tensorboard.delete()\n",
    "except Exception as e:\n",
    "    print(e)\n",
    "\n",
    "if delete_bucket and \"BUCKET_URI\" in globals():\n",
    "    ! gcloud storage rm --recursive $BUCKET_URI"   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "custom_training_tensorboard_profiler.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "python3",
   "name": "tf2-gpu.2-6.m97",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-6:m97"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
